use crate::sync::mutex::Mutex;

use crate::porting::{free, malloc};

use core::alloc::{GlobalAlloc, Layout};
use core::ffi::c_void;
use core::ptr::null_mut;
use libc::realloc;

const MAX_SUPPORTED_ALIGN: usize = 4;

struct Inner {
    alloc: u64,
    max: u64,
}

pub struct AllocatorMonitor {
    mutex: Mutex<Inner>,
}

impl AllocatorMonitor {
    pub fn instance() -> &'static Self {
        unsafe {
            if MON.is_none() {
                MON = Some(Self {
                    mutex: Mutex::new(Inner { alloc: 0, max: 0 }),
                });
            }
            MON.as_ref().unwrap()
        }
    }

    pub fn add_alloc(&self, size: u32) {
        unsafe {
            let mon = MON.as_mut().unwrap();
            let mut inner = mon.mutex.lock().unwrap();
            inner.alloc += size as u64;
            if inner.alloc > inner.max {
                inner.max = inner.alloc;
            }
        }
    }

    pub fn dec_alloc(&self, size: u32) {
        unsafe {
            let mon = MON.as_mut().unwrap();
            let mut inner = mon.mutex.lock().unwrap();
            if size as u64 <= inner.alloc {
                inner.alloc -= size as u64;
            } else {
                panic!("bug");
            }
        }
    }

    pub fn max(&self) -> u64 {
        let inner = self.mutex.lock().unwrap();
        inner.max
    }

    pub fn allocated(&self) -> u64 {
        let inner = self.mutex.lock().unwrap();
        let alloc = inner.alloc;
        alloc
    }
}

static mut MON: Option<AllocatorMonitor> = None;

#[repr(C, align(8))]
pub struct SimpleAllocator {}

unsafe impl Sync for SimpleAllocator {}

unsafe impl GlobalAlloc for SimpleAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let size = layout.size();
        let align = layout.align();
        let _a = if align % MAX_SUPPORTED_ALIGN == 0 {
            align
        } else {
            (align / MAX_SUPPORTED_ALIGN + 1) * MAX_SUPPORTED_ALIGN
        };
        rust_alloc(size) as *mut u8
    }
    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
        rust_free(ptr as *mut c_void);
    }
}

#[cfg(feature = "no_std")]
#[global_allocator]
static ALLOCATOR: SimpleAllocator = SimpleAllocator {};

pub fn set_allocator(
    alloc: unsafe extern "C" fn(libc::size_t) -> *mut c_void,
    fr: unsafe extern "C" fn(*mut c_void),
    realloc: unsafe extern "C" fn(*mut c_void, libc::size_t) -> *mut c_void,
) {
    unsafe {
        MEM_ALLOC = Some(alloc);
        MEM_FREE = Some(fr);
        MEM_REALLOC = Some(realloc);
    }
}

static mut MEM_ALLOC: Option<unsafe extern "C" fn(libc::size_t) -> *mut c_void> = Some(malloc);
static mut MEM_FREE: Option<unsafe extern "C" fn(*mut c_void)> = Some(free);
static mut MEM_REALLOC: Option<unsafe extern "C" fn(*mut c_void, libc::size_t) -> *mut c_void> =
    Some(realloc);

// #[cfg(not(feature = "std"))]
// #[alloc_error_handler]
// fn alloc_error_handler(layout: Layout) -> ! {
//     unsafe {
//         libc::printf(
//             cstr!("allocation error: size %u aligh %u\n\0"),
//             layout.size(),
//             layout.align(),
//         );
//     }
//     panic!("allocate failed!");
// }
//
// #[cfg(not(feature = "std"))]
// #[panic_handler]
// fn panic(_info: &core::panic::PanicInfo) -> ! {
//     loop {}
// }
//
// #[cfg(not(feature = "std"))]
// #[lang = "eh_personality"]
// extern "C" fn eh_personality() {}

// magic: 4 bytes
// 0: 'M'
// 1-3: len, little endian, 24bits
const ALLOC_HEADER: usize = 4;
const MAX_ALLOC_SIZE: usize = 1 << 24;
#[no_mangle]
pub extern "C" fn rust_alloc(size: libc::size_t) -> *mut c_void {
    if size == 0 {
        return null_mut();
    }
    let act_size = size as usize + ALLOC_HEADER;
    if act_size >= MAX_ALLOC_SIZE {
        panic!("alloc {} too big", size);
    }
    let ptr = unsafe { MEM_ALLOC.unwrap()(act_size.into()) as *mut u8 };
    if ptr.is_null() {
        return ptr as *mut c_void;
    }
    let header = unsafe { core::slice::from_raw_parts_mut(ptr, ALLOC_HEADER) };
    header[0] = b'M';
    header[1] = (size & 0xff) as u8;
    header[2] = ((size >> 8) & 0xff) as u8;
    header[3] = ((size >> 16) & 0xff) as u8;
    AllocatorMonitor::instance().add_alloc(act_size as u32);
    unsafe { ptr.add(ALLOC_HEADER) as *mut c_void }
}

#[no_mangle]
pub extern "C" fn rust_realloc(p: *mut c_void, size: libc::size_t) -> *mut c_void {
    if p.is_null() || size == 0 {
        return p;
    }
    let ptr = unsafe { (p as *mut u8).sub(ALLOC_HEADER) };
    let header = unsafe { core::slice::from_raw_parts_mut(ptr, ALLOC_HEADER) };
    if header[0] != b'M' {
        panic!("memory corrupt or point not alloc by rust");
    }
    let old_size = header[1] as usize + ((header[2] as usize) << 8) + ((header[3] as usize) << 16);
    let act_size = size as usize + ALLOC_HEADER;
    if act_size >= MAX_ALLOC_SIZE {
        panic!("alloc {} too big", size);
    }
    let ptr = unsafe { MEM_REALLOC.unwrap()(ptr as *mut c_void, act_size.into()) as *mut u8 };
    if ptr.is_null() {
        return ptr as *mut c_void;
    }
    let header = unsafe { core::slice::from_raw_parts_mut(ptr, ALLOC_HEADER) };
    header[0] = b'M';
    header[1] = (size & 0xff) as u8;
    header[2] = ((size >> 8) & 0xff) as u8;
    header[3] = ((size >> 16) & 0xff) as u8;
    if size >= old_size {
        AllocatorMonitor::instance().add_alloc((size - old_size) as u32);
    } else {
        AllocatorMonitor::instance().dec_alloc((old_size - size) as u32);
    }
    unsafe { ptr.add(ALLOC_HEADER) as *mut c_void }
}

#[no_mangle]
pub extern "C" fn rust_free(p: *mut c_void) {
    if p.is_null() {
        return;
    }
    let ptr = unsafe { (p as *mut u8).sub(ALLOC_HEADER) };
    let header = unsafe { core::slice::from_raw_parts_mut(ptr, ALLOC_HEADER) };
    if header[0] != b'M' {
        panic!("memory corrupt or point not alloc by rust");
    }
    let size = header[1] as usize + ((header[2] as usize) << 8) + ((header[3] as usize) << 16);
    unsafe { MEM_FREE.unwrap()(ptr as *mut c_void) };
    AllocatorMonitor::instance().dec_alloc((size + ALLOC_HEADER) as u32);
}
